/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 *
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#include <vector>
#include <cstdint>
#include <cmath>
#include <random>
#include <iostream>
#include <algorithm>
#include "data_utils.h"
#include "acl/acl.h"
#include "atvc.h"

void AddcmulCustomInt0(uint32_t blockDim, void* stream, uint8_t* tensor1, uint8_t* tensor2, uint8_t* input, uint8_t* output, ATVC::BroadcastParam param, int value);
void AddcmulCustomInt1(uint32_t blockDim, void* stream, uint8_t* tensor1, uint8_t* tensor2, uint8_t* input, uint8_t* output, ATVC::BroadcastParam param, int value);
void AddcmulCustomFloat0(uint32_t blockDim, void* stream, uint8_t* tensor1, uint8_t* tensor2, uint8_t* input, uint8_t* output, ATVC::BroadcastParam param, float value);
void AddcmulCustomFloat1(uint32_t blockDim, void* stream, uint8_t* tensor1, uint8_t* tensor2, uint8_t* input, uint8_t* output, ATVC::BroadcastParam param, float value);

// Addcmul算子的描述：两个输入，一个输出，类型均为float
using BroadcastOpTraits = ATVC::OpTraits<ATVC::OpInputs<float, float, float>, ATVC::OpOutputs<float>, ATVC::OpTemps<float, float, float>>;
using BroadcastOpTraitsInt = ATVC::OpTraits<ATVC::OpInputs<int32_t, int32_t, int32_t>, ATVC::OpOutputs<int32_t>, ATVC::OpTemps<int32_t, int32_t, int32_t>>;

// 负责Broadcast类算子的调度，选择对应的Policy最佳策略并执行Kernel函数
template<class OpTraits>
void BroadcastOpAdapter(uint8_t* tensor1, uint8_t* tensor2, float value, uint8_t* input, uint8_t* output, ATVC::BroadcastParam &param, ATVC::BroadcastPolicy &policy, aclrtStream& stream, bool enableProf=false)
{
    using Inputs = typename OpTraits::In::types;
    using T = typename ATVC::TypeListGet<Inputs, 0>::Type;
    // 申请临时空间workspace，并将其与BroadcastTilingData一同传到Device侧
    uint8_t *workspaceDevice;
    CHECK_ACL(aclrtMalloc((void **)&workspaceDevice, param.workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST));
    param.workspaceAddr = reinterpret_cast<uint64_t>(workspaceDevice);
    // 将tiling api计算出的BroadcastPolicy转化为编译态参数并实例化相应的核函数
        int32_t loopCnt = 1;
    if (enableProf) {
        loopCnt = 20; // 循环20次
    }
    for (int32_t i = 0; i < loopCnt; i++) {
        if constexpr (std::is_same<T, int>::value) {
            if (policy == ATVC::BROADCAST_POLICY0) {
                AddcmulCustomInt0(param.tilingData.coreNum, stream, tensor1, tensor2, input, output, param, value);
            }else if (policy == ATVC::BROADCAST_POLICY1) {
                AddcmulCustomInt1(param.tilingData.coreNum, stream, tensor1, tensor2, input, output, param, value);
            } else {
                printf("[ERROR] Cannot find any matched policy.\n");
            }
        } else {            
            if (policy == ATVC::BROADCAST_POLICY0) {
                AddcmulCustomFloat0(param.tilingData.coreNum, stream, tensor1, tensor2, input, output, param, value);
            }else if (policy == ATVC::BROADCAST_POLICY1) {
                AddcmulCustomFloat1(param.tilingData.coreNum, stream, tensor1, tensor2, input, output, param, value);
            } else {
                printf("[ERROR] Cannot find any matched policy.\n");
            }
        }
    }
    // 流同步后释放申请的param内存
    CHECK_ACL(aclrtSynchronizeStream(stream));
    CHECK_ACL(aclrtFree(workspaceDevice));
}

std::vector<int64_t> GetInputArgs(char* args)
{
    std::string ss(args);
    std::vector<int64_t> vec;
    auto lastPos = ss.find_first_not_of(',', 0);
    auto pos = ss.find_first_of(',', lastPos);
    while(pos != std::string::npos || lastPos != std::string::npos) {
        vec.push_back(std::stoll(ss.substr(lastPos, pos - lastPos)));
        lastPos = ss.find_first_not_of(',', pos);
        pos = ss.find_first_of(',', lastPos);
    }
    return vec;
}

size_t GetShapeSize(std::vector<int64_t> vec) {
    size_t shapeSize = 1;
    for (const auto &s : vec) {
        shapeSize *= s;
    }
    return shapeSize;
}

int32_t main(int32_t argc, char* argv[])
{
 
    std::vector<int64_t> tensor1ShapeIn = GetInputArgs(argv[1]);
    std::vector<int64_t> tensor2ShapeIn = GetInputArgs(argv[2]);
    std::vector<int64_t> valueShapeIn = GetInputArgs(argv[3]);
    std::vector<int64_t> inputShapeIn = GetInputArgs(argv[4]);
    std::vector<int64_t> shapeOut =  GetInputArgs(argv[5]);
    bool intDtype = std::string(argv[6]) == "0";                // 第3个输入，表示算子数据类型
    bool enableProf = std::string(argv[7]) == "1";              // 第4个输入，表示是否使能Profling

    size_t tensor1ByteSize =4 * GetShapeSize(tensor1ShapeIn);
    size_t tensor2ByteSize =4 * GetShapeSize(tensor2ShapeIn);
    size_t inputByteSize =4 * GetShapeSize(inputShapeIn);
    size_t valueByteSize =4 ;
    size_t outputByteSize = 4 * GetShapeSize(shapeOut);
  

    uint8_t* tensor1;
    uint8_t* tensor2;
    uint8_t* value;
    uint8_t* input;
      
    // 初始化Acl资源
    CHECK_ACL(aclInit(nullptr));
    aclrtContext context;
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    CHECK_ACL(aclrtCreateContext(&context, deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));
    uint8_t *outputHost;
    uint8_t *tensor1Device;
    uint8_t *tensor2Device;
    uint8_t *inputDevice;
    uint8_t *outputDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&tensor1), tensor1ByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&tensor2), tensor2ByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&input), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&outputHost), outputByteSize));

    CHECK_ACL(aclrtMalloc((void **)&tensor1Device, tensor1ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&tensor2Device, tensor2ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&inputDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&outputDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_tensor1.bin", tensor1ByteSize, tensor1, tensor1ByteSize);
    ReadFile("./input/input_tensor2.bin", outputByteSize, tensor2, tensor2ByteSize);
    ReadFile("./input/input_input.bin", inputByteSize, input, inputByteSize);

    CHECK_ACL(aclrtMemcpy(tensor1Device, tensor1ByteSize, tensor1, tensor1ByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(tensor2Device, tensor2ByteSize, tensor2, tensor2ByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(inputDevice, inputByteSize, input, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    ATVC::BroadcastParam param;    // Broadcast运行态参数，包含TilingData以及临时空间的相关信息
    ATVC::BroadcastPolicy policy = {-1,-1,-1};  // Broadcast运行态参数，负责映射最适合的Broadcast模板实现
    param.nBufferNum = 1;
    // Host侧调用Tiling API完成相关运行态参数的运算
    if (!ATVC::Host::CalcBroadcastTiling<BroadcastOpTraits>(tensor1ShapeIn, shapeOut, &policy, &param)) {
        printf("Broadcast tiling error.\n");
        return -1;
    };
    if(intDtype){
        // 调用Adapter调度接口，完成核函数的模板调用
        BroadcastOpAdapter<BroadcastOpTraitsInt>(tensor1Device, tensor2Device, 4, inputDevice, outputDevice, param, policy, stream, enableProf);
    }else{
        BroadcastOpAdapter<BroadcastOpTraits>(tensor1Device, tensor2Device, 4.0f, inputDevice, outputDevice, param, policy, stream, enableProf);
    }

    CHECK_ACL(aclrtMemcpy(outputHost, outputByteSize, outputDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    std::vector<float> output(reinterpret_cast<float*>(outputHost), reinterpret_cast<float*>(outputHost) + outputByteSize);
    WriteFile("./output/output_output.bin", outputHost, outputByteSize);

    // 释放Acl资源
    CHECK_ACL(aclrtFree(tensor1Device));
    CHECK_ACL(aclrtFree(tensor2Device));
    CHECK_ACL(aclrtFree(inputDevice));
    CHECK_ACL(aclrtFree(outputDevice));
    CHECK_ACL(aclrtFreeHost(outputHost));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtDestroyContext(context));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize()); 
    printf("Accuracy verification passed.\n");
    return 0;
}
